from fastapi import APIRouter, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
import os
import shutil
import uuid
from typing import Optional
from paddleocrPdfGpuDeal import extract_text_from_pdf
# 导入配置
from config import Config

router = APIRouter()

@router.post("/process_pdf")
async def process_pdf(
    file: UploadFile = File(...),
    watermark_removal_method: Optional[str] = Form("color_filter"),
    lang: Optional[str] = Form("ch")
):
    """
    使用CPU处理上传的PDF文件并返回OCR识别结果
    
    参数:
        file: 上传的PDF文件
        watermark_removal_method: 去水印方法，可选 'basic', 'advanced', 'color_filter', 'none'
        lang: 识别语言，默认为中文
    
    返回:
        JSON格式的OCR识别结果
    """
    # 检查文件类型
    if file.filename is None or not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="只接受PDF文件")
    
    try:
        # 创建临时目录
        temp_dir = os.path.join(Config.TEMP_DIR, str(uuid.uuid4()))  # 使用配置文件中的临时目录
        os.makedirs(temp_dir, exist_ok=True)
        
        # 保存上传的PDF文件
        temp_pdf_path = os.path.join(temp_dir, "uploaded.pdf")
        with open(temp_pdf_path, "wb") as pdf_file:
            pdf_file.write(await file.read())
        
        # 调用CPU OCR处理函数
        result = extract_text_from_pdf(
            temp_pdf_path, 
            temp_output_dir=temp_dir, 
            watermark_removal_method=watermark_removal_method or "color_filter",
            lang=lang or "ch"
        )
        
        # 构建JSON响应
        response_data = {
            "filename": file.filename,
            "pages_count": len(result),
            "ocr_result": result
        }
        
        # 清理临时文件
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        
        return JSONResponse(content=response_data)
    
    except Exception as e:
        # 确保清理临时文件
        try:
            # 使用try-except确保temp_dir变量存在且不为None
            temp_dir_local = locals().get('temp_dir')
            if temp_dir_local and os.path.exists(temp_dir_local):
                shutil.rmtree(temp_dir_local)
        except:
            pass  # 忽略清理错误
        
        # 返回错误信息
        raise HTTPException(status_code=500, detail=f"处理PDF时出错: {str(e)}")